from util import *

def extract_first_n_rows(sub_dataset_dir,n=1000):
    """
    从输入的CSV文件中提取前n行，并将结果保存为新的CSV文件。

    参数:
    sub_dataset_dir:。
    n (int): 提取的行数，默认为1000。
    """
    
    file_list=get_file_list_in_dir(sub_dataset_dir)
    file_list.sort()
    input_file=os.path.join(sub_dataset_dir,file_list[0])
    output_file=os.path.join(sub_dataset_dir,file_list[0]+".csv")
    
    # 读取CSV文件
    data = read_csv_from_tar_gz(input_file)

    # 提取前n行
    extracted_data = data.head(n)

    # 保存为新的CSV文件
    extracted_data.to_csv(output_file, index=False)
    print(f"成功提取前{n}行并保存到 {output_file}。")

    


def generate_view_file():
    root_dir=get_root_dir()
    dataset_dir=os.path.dirname(root_dir)+"/MS_data/data"
    dir="CallGraph"
    extract_first_n_rows(os.path.join(dataset_dir, dir))
    dir = "MSMetrics"
    extract_first_n_rows(os.path.join(dataset_dir, dir))
    dir = "MSRTMCR"
    extract_first_n_rows(os.path.join(dataset_dir, dir))
    dir = "NodeMetrics"
    extract_first_n_rows(os.path.join(dataset_dir, dir))
    
generate_view_file()